"""
IQL from demo experiment

nohup python antmaze_finetune.py > data/gogogo.txt 2>&1 &
nohup python antmaze_finetune.py > /dev/null 2>&1 &

CUDA_VISIBLE_DEVICES=1; nohup python antmaze_finetune.py > /dev/null 2>&1 &

"""

import os
import sys
import argparse
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

from rlkit.demos.source.hdf5_path_loader import HDF5PathLoader
from rlkit.launchers.experiments.awac.finetune_rl import experiment, process_args

from rlkit.launchers.launcher_util import run_experiment

from rlkit.torch.sac.policies import GaussianPolicy
from rlkit.torch.sac.iql_trainer import IQLTrainer

import random

import d4rl



def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--dataset_path', type=str, 
                       default='/home/ubunut/KXT/code/datasets/offline_rl/ant_maze_new/Ant_maze_big-maze_noisy_multistart_True_multigoal_True_sparse.hdf5',
                       help='Path to the dataset')
    parser.add_argument('--remark', type=str, default='', help='remark')
    parser.add_argument('--dataset_length', type=int, default=10000, help='Length of dataset to use')
    parser.add_argument('--env_id', type=str, default='antmaze-medium-diverse-v0', help='Environment ID')
    parser.add_argument('--seed', type=int, default=0, help='Random seed')
    # jsrl
    """
    max_guide_steps: 计算得出，完成任务的平均步数
    curriculum_n: 从1到guide_steps中尝试，较小的值即可
    curriculum_threshold: 进入下一个阶段的阈值（容差度）
    """
    parser.add_argument('--dobule_buffer', type=int, default=1, help='是否使用双缓冲区，0表示不使用，1表示使用')
    parser.add_argument('--guide_mode', type=str, default=None, help='jsrl, pretrain, self-jsrl')
    parser.add_argument('--guide_policy_path', type=str, default=None)
    parser.add_argument('--max_guide_steps', type=int, default=400)
    parser.add_argument('--curriculum_n', type=int, default=20)
    parser.add_argument('--curriculum_threshold', type=float, default=0.95)
    return parser.parse_args()


def get_variant(args):

    return dict(
        # 训练参数
        algo_kwargs=dict(
            start_epoch=-1000, # offline epochs
            num_epochs=1001, # online epochs
            batch_size=256,
            eval_freq=20, # 评估频率
            num_eval_steps_per_epoch=10000,  # 评估步数（最少评估10轮，采满10000步）
            num_trains_per_train_loop=1000,  # 每个内部训练循环（train_loop）中，参数更新的次数。
            num_expl_steps_per_train_loop=1000,  # 每个内部训练循环中，智能体在环境中探索的步数。（最大环境步数）
            min_num_steps_before_training=5000,  # 在正式训练开始前，需要预先收集的探索步数。
        ),
        # 测试参数
        # algo_kwargs=dict(
        #     start_epoch=-10, # offline epochs
        #     num_epochs=11, # online epochs
        #     batch_size=256,
        #     eval_freq=5, # 评估频率
        #     num_eval_steps_per_epoch=1000,
        #     num_trains_per_train_loop=1,
        #     num_expl_steps_per_train_loop=1000,
        #     min_num_steps_before_training=256,
        # ),
        jsrl_kwargs=dict(
            guide_policy_path=args.guide_policy_path,  # jsrl 指导策略路径
            max_guide_steps=args.max_guide_steps,  # jsrl 指导步数
            curriculum_n=args.curriculum_n,  # jsrl 课程学习阶段数
            curriculum_threshold=args.curriculum_threshold,  # jsrl 课程学习阶段阈值
        ),
        max_path_length=1000,
        replay_buffer_size=int(1E6),  # 1M 数据集最大长度
        expl_replay_buffer_size=int(1E5),  # jsrl 配置
        dobule_buffer=args.dobule_buffer,  # 是否使用双缓冲区
        layer_size=256,
        policy_class=GaussianPolicy,
        policy_kwargs=dict(
            hidden_sizes=[256, 256, ],
            max_log_std=0,
            min_log_std=-6,
            std_architecture="values",
        ),
        qf_kwargs=dict(
            hidden_sizes=[256, 256, ],
        ),

        version="normal",
        collection_mode='batch',
        trainer_class=IQLTrainer,
        trainer_kwargs=dict(
            discount=0.99,
            policy_lr=3E-4,
            qf_lr=3E-4,
            reward_scale=1,

            policy_weight_decay=0,
            q_weight_decay=0,

            reward_transform_kwargs=dict(m=1, b=-1),
            terminal_transform_kwargs=None,

            beta=0.1,
            quantile=0.9,
            clip_score=100,
        ),
        launcher_config=dict(
            num_exps_per_instance=1,
            region='us-west-2',
        ),

        path_loader_class=HDF5PathLoader,
        path_loader_kwargs=dict(),
        add_env_demos=False,
        add_env_offpolicy_data=False,

        load_demos=False,
        load_env_dataset_demos=True,
        dataset_path=args.dataset_path,
        dataset_length=args.dataset_length,

        normalize_env=False,
        env_id=args.env_id,

        seed=args.seed,
    )

def main():
    args = parse_args()
    variant = get_variant(args)

    # jsrl配置
    # TODO 热启动
    if args.guide_mode == 'jsrl':
        variant['algo_kwargs']['start_epoch'] = 0
        variant['algo_kwargs']['eval_freq'] = 10
        variant['dobule_buffer'] = 1
    elif args.guide_mode == 'pretrain':
        variant['algo_kwargs']['num_epochs'] = 0
        variant['dobule_buffer'] = 0
    elif args.guide_mode == 'self':
        variant['algo_kwargs']['eval_freq'] = 10
        variant['dobule_buffer'] = 1

    remark = ''
    if args.guide_mode == 'jsrl':
        remark += f'-jsrl-{args.max_guide_steps}-n{args.curriculum_n}-th{args.curriculum_threshold}'
    elif args.guide_mode == 'self':
        remark += f'-self-{args.max_guide_steps}-n{args.curriculum_n}-th{args.curriculum_threshold}'
    elif args.guide_mode == 'pretrain':
        remark += '-pretrain'
    if variant['dobule_buffer']:
        print('Using double buffer')
        remark += '-dbuf'
    if args.remark:
        remark += f'-{args.remark}'
        
    run_experiment(experiment,
        variant=variant,
        exp_prefix=f'iql-{args.env_id}-{args.dataset_length // 1000}k{remark}',
        # exp_prefix='iql-antmaze-medium-diverse-v0-10k',
        mode="here_no_doodad",
        unpack_variant=False,
        use_gpu=True,
    )

if __name__ == "__main__":
    main()
